#include "stdafx.h"

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <tchar.h>
#include <WinIoCtl.h>
#include <ntddscsi.h>

#include "SMARTDrive.h"

using namespace System::Runtime::InteropServices;

namespace MaxSMART {

#pragma pack(1)
	struct IDENTIFY_DATA {
		USHORT GeneralConfiguration;            // 00 00
		USHORT NumberOfCylinders;               // 02  1
		USHORT Reserved1;                       // 04  2
		USHORT NumberOfHeads;                   // 06  3
		USHORT UnformattedBytesPerTrack;        // 08  4
		USHORT UnformattedBytesPerSector;       // 0A  5
		USHORT SectorsPerTrack;                 // 0C  6
		USHORT VendorUnique1[3];                // 0E  7-9
		USHORT SerialNumber[10];                // 14  10-19
		USHORT BufferType;                      // 28  20
		USHORT BufferSectorSize;                // 2A  21
		USHORT NumberOfEccBytes;                // 2C  22
		USHORT FirmwareRevision[4];             // 2E  23-26
		USHORT ModelNumber[20];                 // 36  27-46
		UCHAR  MaximumBlockTransfer;            // 5E  47
		UCHAR  VendorUnique2;                   // 5F
		USHORT DoubleWordIo;                    // 60  48
		USHORT Capabilities;                    // 62  49
		USHORT Reserved2;                       // 64  50
		UCHAR  VendorUnique3;                   // 66  51
		UCHAR  PioCycleTimingMode;              // 67
		UCHAR  VendorUnique4;                   // 68  52
		UCHAR  DmaCycleTimingMode;              // 69
		USHORT TranslationFieldsValid:1;        // 6A  53
		USHORT Reserved3:15;
		USHORT NumberOfCurrentCylinders;        // 6C  54
		USHORT NumberOfCurrentHeads;            // 6E  55
		USHORT CurrentSectorsPerTrack;          // 70  56
		ULONG  CurrentSectorCapacity;           // 72  57-58
		USHORT CurrentMultiSectorSetting;       //     59
		ULONG  UserAddressableSectors;          //     60-61
		USHORT SingleWordDMASupport : 8;        //     62
		USHORT SingleWordDMAActive : 8;
		USHORT MultiWordDMASupport : 8;         //     63
		USHORT MultiWordDMAActive : 8;
		USHORT AdvancedPIOModes : 8;            //     64
		USHORT Reserved4 : 8;
		USHORT MinimumMWXferCycleTime;          //     65
		USHORT RecommendedMWXferCycleTime;      //     66
		USHORT MinimumPIOCycleTime;             //     67
		USHORT MinimumPIOCycleTimeIORDY;        //     68
		USHORT Reserved5[2];                    //     69-70
		USHORT ReleaseTimeOverlapped;           //     71
		USHORT ReleaseTimeServiceCommand;       //     72
		USHORT MajorRevision;                   //     73
		USHORT MinorRevision;                   //     74
		USHORT Reserved6[50];                   //     75-126
		USHORT SpecialFunctionsEnabled;         //     127
		USHORT Reserved7[128];                  //     128-255
	};
#pragma pack()

array<SmartDrive^>^ SmartDrive::EnumeratePhysicalDrives()
{
	List<SmartDrive^>^ ret = gcnew List<SmartDrive^>();

	for (unsigned char drive = 0; drive < 32; drive++)
	{
//		HANDLE handle = 0;
		TCHAR name[MAX_PATH];
		_stprintf_s(name, _T("\\\\.\\PhysicalDrive%d\\"), drive); 
 		if (GetDriveType(name) != DRIVE_FIXED) 
 			continue; // skip removable disks
 		else
			name[_tcslen(name) - 1] = 0; // removing the trailing backslash
			
		SafeFileHandle^ handle = gcnew SafeFileHandle((IntPtr) CreateFile (name, GENERIC_READ | GENERIC_WRITE, FILE_SHARE_DELETE | FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL), true);
		if (!handle->IsInvalid) {
			ret->Add(gcnew SmartDrive(handle, Marshal::PtrToStringAnsi((IntPtr) name), true));
		} else {
			// On Vista/7 with no rights ?
			handle = gcnew SafeFileHandle((IntPtr) CreateFile (name, 0, FILE_SHARE_DELETE | FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL), true);
			if (!handle->IsInvalid) {
				ret->Add(gcnew SmartDrive(handle, Marshal::PtrToStringAnsi((IntPtr) name), false));
			}
		}
	}
	return ret->ToArray() ;
}

union STORAGE_DEVICE_DESCRIPTOR_DATA {
	STORAGE_DEVICE_DESCRIPTOR desc;
	char raw[256];
};

SmartDrive::SmartDrive(SafeFileHandle^ handle, String^ name, bool hasRights)
	: Handle(handle), _name(name), _hasRights(hasRights),
	_attributes(gcnew Dictionary<SmartAttribute, SmartAttributeData>()),
	_thresholds(gcnew Dictionary<SmartAttribute, Byte>())
{
	MaxSMART::BusType busType;
	if (ExecuteStorageQueryProperty(busType)) _busType = busType;

	SCSI_ADDRESS address = {};
	address.Length = sizeof(SCSI_ADDRESS);

	DWORD bytes = 0;
	if (FALSE != DeviceIoControl(Handle->DangerousGetHandle().ToPointer(), IOCTL_SCSI_GET_ADDRESS, NULL, 0,  &address, sizeof(address), &bytes, NULL)) {
		this->_scsiAddress = gcnew MaxSMART::ScsiAddress(address.PortNumber, address.PathId, address.TargetId, address.Lun);;
 		if ((BusType == MaxSMART::BusType::Ata) || (BusType == MaxSMART::BusType::Sata)) {
 			Index = (address.PathId * 2) + address.TargetId; // http://www.microsoft.com/whdc/device/storage/ATAPortInfo.mspx
 		}
	}
}

SmartDrive::~SmartDrive() {
	if (Handle != nullptr) Handle->Close();
}

bool SmartDrive::ExecuteStorageQueryProperty(MaxSMART::BusType% busType) {
	bool ret = false;

	STORAGE_PROPERTY_QUERY query = {};
	query.PropertyId = StorageDeviceProperty;
	query.QueryType = PropertyStandardQuery;
	STORAGE_DEVICE_DESCRIPTOR_DATA data = {};

	DWORD bytes = 0;
	if (FALSE != DeviceIoControl(Handle->DangerousGetHandle().ToPointer(), IOCTL_STORAGE_QUERY_PROPERTY, &query, sizeof(query), &data, sizeof(data), &bytes, NULL)) {
		busType = (MaxSMART::BusType) data.desc.BusType;
		if (String::IsNullOrEmpty(Model)) {
			Model = Marshal::PtrToStringAnsi((IntPtr) (data.desc.VendorIdOffset ? data.raw + data.desc.VendorIdOffset : ""))->Trim() +
				Marshal::PtrToStringAnsi((IntPtr) (data.desc.ProductIdOffset ? data.raw + data.desc.ProductIdOffset : ""))->Trim();
		}
		if (String::IsNullOrEmpty(Serial) && (data.desc.SerialNumberOffset != 0) && (data.desc.SerialNumberOffset != -1)) {
			CHAR serial[21] = {};
			if (strlen(data.raw + data.desc.SerialNumberOffset) > 20) { // nasty bug of some drivers or what ?
				CHAR buffer[3] = "00";
				int chr;
				for(unsigned int i = 0; i < strlen(data.raw + data.desc.SerialNumberOffset) / 2; i++) {
					buffer[0] = data.raw[data.desc.SerialNumberOffset + i * 2];
					buffer[1] = data.raw[data.desc.SerialNumberOffset + i * 2 + 1];
					sscanf_s(buffer, "%x", &chr);
					serial[i] = chr;
				}
				for(int i = 0; i < 10; i++) {
					char tmp = serial[i * 2];
					serial[i * 2] = serial[i * 2 + 1];
					serial[i * 2 + 1] = tmp;
				}
			} else strcpy_s(serial, data.raw + data.desc.SerialNumberOffset); // sometimes you need _swab but when ?
			Serial = Marshal::PtrToStringAnsi((IntPtr) serial)->Trim();
		}
		ret = true;
	}

	return ret;
}

DWORD GetBufferSize(SmartCommand command) {
	DWORD ret = 0;
	switch (command) {
	case SmartCommand::IdentifyDevice:
		ret = IDENTIFY_BUFFER_SIZE; break;
	case SmartCommand::ReadAttributes:
		ret = READ_ATTRIBUTE_BUFFER_SIZE; break;
	case SmartCommand::ReadThresholds:
		ret = READ_THRESHOLD_BUFFER_SIZE; break;
	}
	return ret;
}

void FillSendCmdInParams(PSENDCMDINPARAMS params, BYTE drive, SmartCommand command, bool atapi = false) {
	assert(params != NULL);
	params->cBufferSize = GetBufferSize(command);
	params->bDriveNumber = drive;
	params->irDriveRegs.bFeaturesReg = 0;
	switch (command) {
	case SmartCommand::EnableSmart:
		params->irDriveRegs.bFeaturesReg = ENABLE_SMART; break;
	case SmartCommand::DisableSmart:
		params->irDriveRegs.bFeaturesReg = DISABLE_SMART; break;
	case SmartCommand::ReadAttributes:
		params->irDriveRegs.bFeaturesReg = READ_ATTRIBUTES; break;
	case SmartCommand::ReadThresholds:
		params->irDriveRegs.bFeaturesReg = READ_THRESHOLDS; break;
	}
	params->irDriveRegs.bSectorCountReg = 1;
	params->irDriveRegs.bSectorNumberReg = 1;
	params->irDriveRegs.bCylLowReg = SMART_CYL_LOW;
	params->irDriveRegs.bCylHighReg = SMART_CYL_HI;
	params->irDriveRegs.bDriveHeadReg = (drive & 1) ? 0xB0 : 0xA0; // MASTER = A0, SLAVE = B0.
	params->irDriveRegs.bCommandReg = (command == SmartCommand::IdentifyDevice) ? (atapi ? ATAPI_ID_CMD : ID_CMD) : SMART_CMD;
}

void SmartDrive::UpdateIdentifyData(PIDENTIFY_DATA data) {
	assert(data != NULL);
	if ((String::IsNullOrEmpty(Model)) && (data->ModelNumber[0] != 0)) {
		CHAR model[sizeof(data->ModelNumber) + 1] = {};
		_swab((PCHAR) data->ModelNumber, model, sizeof(data->ModelNumber));
		Model = Marshal::PtrToStringAnsi((IntPtr) model)->Trim();
	}
	if ((String::IsNullOrEmpty(Serial)) && (data->SerialNumber[0] != 0)) {
		CHAR serial[sizeof(data->SerialNumber) + 1] = {};
		_swab((PCHAR) data->SerialNumber, serial, sizeof(data->SerialNumber));
		Serial = Marshal::PtrToStringAnsi((IntPtr) serial)->Trim();
	}
}

bool SmartDrive::ExecuteSmart(SmartCommand command)
{
	bool ret = false;

	HANDLE handle = Handle->DangerousGetHandle().ToPointer();

	GETVERSIONINPARAMS GetVersionParams = {};
	DWORD bytes = 0;
	if (FALSE != DeviceIoControl(handle, SMART_GET_VERSION, NULL, 0, &GetVersionParams, sizeof(GETVERSIONINPARAMS), &bytes, NULL)) {
		SENDCMDINPARAMS inParams = {};
		DWORD bufferSize = GetBufferSize(command);
		DWORD outSize = sizeof(SENDCMDOUTPARAMS) + ((bufferSize == 0) ? 0 : (bufferSize - 1));
		PBYTE outBuffer = new BYTE[outSize];
		FillSendCmdInParams(&inParams, Index, command, ((GetVersionParams.bIDEDeviceMap >> Index) & 0x10) != 0);
		DWORD code;
		switch (command) {
		case SmartCommand::EnableSmart:
		case SmartCommand::DisableSmart:
			code = SMART_SEND_DRIVE_COMMAND; break;
		default:
			code = SMART_RCV_DRIVE_DATA; break;
		}
		if (FALSE != DeviceIoControl(handle, code, &inParams, sizeof(inParams), outBuffer, outSize, &bytes, NULL)) {
			if (command == SmartCommand::IdentifyDevice) UpdateIdentifyData((PIDENTIFY_DATA) ((PSENDCMDOUTPARAMS) outBuffer)->bBuffer);
			else if (command == SmartCommand::ReadAttributes) UpdateAttributes(((PSENDCMDOUTPARAMS) outBuffer)->bBuffer, bufferSize);
			else if (command == SmartCommand::ReadThresholds) ret = UpdateThresholds(((PSENDCMDOUTPARAMS) outBuffer)->bBuffer, bufferSize);
			ret = true;
		}
		delete[] outBuffer;
	}

	return ret;
}

//#define IOCTL_SCSI_MINIPORT 0x0004D008
// 
// typedef struct _SRB_IO_CONTROL { 
// 	ULONG  HeaderLength; 
// 	UCHAR  Signature[8]; 
// 	ULONG  Timeout; 
// 	ULONG  ControlCode; 
// 	ULONG  ReturnCode; 
// 	ULONG  Length; 
// } SRB_IO_CONTROL, *PSRB_IO_CONTROL;
// 
//#define FILE_DEVICE_SCSI 0x0000001b 
#define IOCTL_SCSI_MINIPORT_SMART_VERSION           ((FILE_DEVICE_SCSI << 16) + 0x0500)
#define IOCTL_SCSI_MINIPORT_IDENTIFY                ((FILE_DEVICE_SCSI << 16) + 0x0501)
#define IOCTL_SCSI_MINIPORT_READ_SMART_ATTRIBS      ((FILE_DEVICE_SCSI << 16) + 0x0502)
#define IOCTL_SCSI_MINIPORT_READ_SMART_THRESHOLDS   ((FILE_DEVICE_SCSI << 16) + 0x0503)
#define IOCTL_SCSI_MINIPORT_ENABLE_SMART            ((FILE_DEVICE_SCSI << 16) + 0x0504)
#define IOCTL_SCSI_MINIPORT_DISABLE_SMART           ((FILE_DEVICE_SCSI << 16) + 0x0505)
#define IOCTL_SCSI_MINIPORT_RETURN_STATUS           ((FILE_DEVICE_SCSI << 16) + 0x0506)

typedef struct _SENDCMDINOUTPARAMS {
	SRB_IO_CONTROL srb;
	union {
		SENDCMDINPARAMS in;
		SENDCMDOUTPARAMS out;
	} params;
} SRB_SENDCMDINOUTPARAMS, *PSRB_SENDCMDINOUTPARAMS;

bool SmartDrive::ExecuteScsiMiniport(SmartCommand command) {
	bool ret = false;

	HANDLE handle = Handle->DangerousGetHandle().ToPointer();

	DWORD bufferSize = GetBufferSize(command);
	DWORD outSize = sizeof(SRB_SENDCMDINOUTPARAMS) + ((bufferSize == 0) ? 0 : (bufferSize - 1)) ;
	PBYTE outBuffer = new BYTE[outSize];
	PSRB_SENDCMDINOUTPARAMS inBuffer =  (PSRB_SENDCMDINOUTPARAMS) outBuffer;

	memset(outBuffer, 0, outSize);
	inBuffer->srb.HeaderLength = sizeof(SRB_IO_CONTROL);
	inBuffer->srb.Timeout = 10;
	inBuffer->srb.Length = outSize - sizeof(SRB_IO_CONTROL);
	inBuffer->srb.ControlCode = 0;
	switch(command) {
	case SmartCommand::IdentifyDevice:
		inBuffer->srb.ControlCode = IOCTL_SCSI_MINIPORT_IDENTIFY; break;
	case SmartCommand::EnableSmart:
		inBuffer->srb.ControlCode = IOCTL_SCSI_MINIPORT_ENABLE_SMART; break;
	case SmartCommand::DisableSmart:
		inBuffer->srb.ControlCode = IOCTL_SCSI_MINIPORT_DISABLE_SMART; break;
	case SmartCommand::ReadAttributes:
		inBuffer->srb.ControlCode = IOCTL_SCSI_MINIPORT_READ_SMART_ATTRIBS; break;
	case SmartCommand::ReadThresholds:
		inBuffer->srb.ControlCode = IOCTL_SCSI_MINIPORT_READ_SMART_THRESHOLDS; break;
	}
	memcpy_s(inBuffer->srb.Signature, sizeof(inBuffer->srb.Signature), "SCSIDISK", 8);
	FillSendCmdInParams(&inBuffer->params.in, Index, command);

	DWORD bytes = 0;
	if (FALSE != DeviceIoControl(handle, IOCTL_SCSI_MINIPORT, inBuffer, sizeof(SRB_IO_CONTROL) + sizeof(SENDCMDINPARAMS), outBuffer, outSize, &bytes, NULL)) {
		if (command == SmartCommand::IdentifyDevice) UpdateIdentifyData((PIDENTIFY_DATA) ((PSRB_SENDCMDINOUTPARAMS) outBuffer)->params.out.bBuffer);
		else if (command == SmartCommand::ReadAttributes) ret = UpdateAttributes(((PSRB_SENDCMDINOUTPARAMS) outBuffer)->params.out.bBuffer, bufferSize);
		else if (command == SmartCommand::ReadThresholds) ret = UpdateThresholds(((PSRB_SENDCMDINOUTPARAMS) outBuffer)->params.out.bBuffer, bufferSize);
	}
	delete[] outBuffer;

	return ret;
}

bool SmartDrive::ExecuteStoragePredictFailure()
{
	bool ret = false;

	STORAGE_PREDICT_FAILURE outBuffer = {};
	DWORD bytes = 0;
	if (FALSE != DeviceIoControl(Handle->DangerousGetHandle().ToPointer(), IOCTL_STORAGE_PREDICT_FAILURE, NULL, 0, &outBuffer, sizeof(outBuffer), &bytes, NULL)) {
		ret = UpdateAttributes(outBuffer.VendorSpecific, sizeof(outBuffer.VendorSpecific));
	}

	return ret;
}


bool SmartDrive::UpdateAttributes( PBYTE buffer, DWORD length )
{
	if ((buffer == NULL) || (length < READ_ATTRIBUTE_BUFFER_SIZE)) return false;

	for(int i = 0; i < 30; i++) {
		int offset = 2 + i * 12;
		if(buffer[offset] != 0) {
			SmartAttribute attribute = (SmartAttribute) buffer[offset];
			SmartAttributeData data;
			data.Flags = (SmartAttributeFlags) (buffer[offset + 1] + (buffer[offset + 2] << 8));
			data.Value = buffer[offset + 3];
			data.Peak = buffer[offset + 4];
			data.Raw1 = buffer[offset + 5] + (buffer[offset + 6] << 8);
			data.Raw2 = buffer[offset + 7] + (buffer[offset + 8] << 8);
			data.Raw3 = buffer[offset + 9] + (buffer[offset + 10] << 8);
			data.Reserved = buffer[offset + 11];
			Attributes[attribute] = data;
		}
	}

	return true;
}

bool SmartDrive::UpdateThresholds( PBYTE buffer, DWORD length )
{
	if ((buffer == NULL) || (length < READ_THRESHOLD_BUFFER_SIZE)) return false;

	for(int i = 0; i < 30; i++) {
		int offset = 2 + i * 12;
		if(buffer[offset] != 0) {
			SmartAttribute attribute = (SmartAttribute) buffer[offset];
			Thresholds[attribute] = buffer[offset + 1];
		}
	}

	return true;
}

bool SmartDrive::ExecuteCommand( SmartMethod method, SmartCommand command )
{
	switch (method) {
	case SmartMethod::Smart:
		return ExecuteSmart(command);
		break;
	case SmartMethod::ScsiMiniport:
		return ExecuteScsiMiniport(command);
		break;
	case SmartMethod::Storage:
		switch(command){
		case SmartCommand::IdentifyDevice:
			MaxSMART::BusType busType;
			return ExecuteStorageQueryProperty(busType);
			break;
		case SmartCommand::ReadAttributes:
			return ExecuteStoragePredictFailure();
			break;
		default:
			throw gcnew System::ArgumentException();
		}
		break;
	default:
		throw gcnew System::ArgumentException();
	}
}

}
